# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
r"""Downloads and converts Flowers data to TFRecords of TF-Example protos.

This module downloads the Flowers data, uncompresses it, reads the files
that make up the Flowers data and creates two TFRecord datasets: one for train
and one for test. Each TFRecord dataset is comprised of a set of TF-Example
protocol buffers, each of which contain a single image and label.

The script should take about a minute to run.

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
import os
import random
import sys

import tensorflow as tf
import pandas as pd
import numpy as np

from datasets import dataset_utils

_NUM_VALIDATION = 7000

# Seed for repeatability.
_RANDOM_SEED = 0

# The number of shards per dataset split.
_NUM_SHARDS = 2


class ImageReader(object):
    """Helper class that provides TensorFlow image coding utilities."""

    def __init__(self):
        # Initializes function that decodes RGB JPEG data.
        self._decode_jpeg_data = tf.placeholder(dtype=tf.string)
        # tf.image.decode_jpeg解码jpg格式图片
        self._decode_jpeg = tf.image.decode_jpeg(self._decode_jpeg_data, channels=3)
        self._decode_jpeg = tf.image.resize_images(self._decode_jpeg, (224, 224))

    def read_image_dims(self, sess, image_data):
        # image = self.decode_jpeg(sess, image_data)
        # return image.shape[0], image.shape[1]
        return 224,224

    def decode_jpeg(self, sess, image_data):
        image = sess.run(self._decode_jpeg,
                         feed_dict={self._decode_jpeg_data: image_data})
        assert len(image.shape) == 3
        assert image.shape[2] == 3
        return image


def _get_filenames_and_classes(dataset_dir):
    """Returns a list of filenames and inferred class names.

    Args:
      dataset_dir: A directory containing a set of subdirectories representing
        class names. Each subdirectory should contain PNG or JPG encoded images.

    Returns:
      A list of image file paths, relative to `dataset_dir` and the list of
      subdirectories, representing class names.
    """
    # flower_root是数据集的路径
    flower_root = os.path.join(dataset_dir, 'mydata')
    directories = []
    class_names = []
    print(flower_root)
    # 列出数据集所在路径的所有文件夹
    for filename in os.listdir(flower_root):
        print(filename)
        # 这些文件夹的所在路径
        path = os.path.join(flower_root, filename)
        # 检验路径是否为目录，若是，将路径和目录名分别放入directories和class_names
        if os.path.isdir(path):
            directories.append(path)
            class_names.append(filename)

    # 收集所有图片的存储路径，放入photo_filenames
    photo_filenames = []
    for directory in directories:
        for filename in os.listdir(directory):
            path = os.path.join(directory, filename)
            photo_filenames.append(path)

    # photo_filenames是所有图片的存储路径，class_names是分类名(不同类图片的文件夹名)
    return photo_filenames, sorted(class_names)


def _get_dataset_filename(dataset_dir, split_name, shard_id):
    # 编写输出文件的描述形式
    output_filename = 'mydata_%s_%05d-of-%05d.tfrecord' % (
        split_name, shard_id, _NUM_SHARDS)
    return os.path.join(dataset_dir, output_filename)


def _convert_dataset(split_name, filenames, labels, dataset_dir):
    """Converts the given filenames to a TFRecord dataset.

    Args:
      split_name: The name of the dataset, either 'train' or 'validation'.
      filenames: A list of absolute paths to png or jpg images.
      class_names_to_ids: A dictionary from class names (strings) to ids
        (integers).
      dataset_dir: The directory where the converted datasets are stored.
    """
    assert split_name in ['train', 'validation']

    # 设定每个shard包含的图片数目
    num_per_shard = int(math.ceil(len(filenames) / float(_NUM_SHARDS)))

    with tf.Graph().as_default():
        image_reader = ImageReader()

        with tf.Session('') as sess:

            for shard_id in range(_NUM_SHARDS):
                # 编写输出文件的描述形式
                output_filename = _get_dataset_filename(
                    dataset_dir, split_name, shard_id)

                with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
                    start_ndx = shard_id * num_per_shard
                    end_ndx = min((shard_id + 1) * num_per_shard, len(filenames))
                    for i in range(start_ndx, end_ndx):
                        sys.stdout.write('\r>> Converting image %d/%d shard %d' % (
                            i + 1, len(filenames), shard_id))
                        sys.stdout.flush()

                        # Read the filename:
                        # tf.gfile.FastGFile('path','rb').read()读取没有经过解码的图片
                        image_data = tf.gfile.FastGFile(filenames[i], 'rb').read()
                        # 返回图片的长宽
                        height, width = image_reader.read_image_dims(sess, image_data)

                        # os.path.dirname去掉文件名，返回目录 
                        # os.path.basename返回目录最后的文件名
                        # class_name就是图片所在的文件夹名
                        # class_name = os.path.basename(os.path.dirname(filenames[i]))
                        # class_id = class_names_to_ids[class_name]
                        class_id = list(labels[i])

                        example = dataset_utils.image_to_tfexample(
                            image_data, b'jpg', height, width, class_id)
                        tfrecord_writer.write(example.SerializeToString())

    sys.stdout.write('\n')
    sys.stdout.flush()


def _dataset_exists(dataset_dir):
    for split_name in ['train', 'validation']:
        for shard_id in range(_NUM_SHARDS):
            # _get_dataset_filename返回输出文件的路径和名字的结合
            output_filename = _get_dataset_filename(
                dataset_dir, split_name, shard_id)
            # 如果相应路径的输出文件已经存在，返回False
            if not tf.gfile.Exists(output_filename):
                return False
    return True
	
	

def run(dataset_dir):
    """Runs the download and conversion operation.

    Args:
        dataset_dir: The dataset directory where the dataset is stored.
    """
    # 检测数据集的路径是否存在，如果不存在就创建新的
    if not tf.gfile.Exists(dataset_dir):
        tf.gfile.MakeDirs(dataset_dir)
    # 如果相应路径的输出文件已经存在，返回相应提示
    if _dataset_exists(dataset_dir):
        print('Dataset files already exist. Exiting without re-creating them.')
        return

    # photo_filenames包含所有图片的存储路径，class_names是所有的分类名
    # photo_filenames, class_names = _get_filenames_and_classes(dataset_dir)
    csv_path = os.path.join(dataset_dir, 'visual_china_train.csv')
    photo_filenames = pd.read_csv(csv_path, sep=',', usecols=['img_path'], encoding='utf-8')
    photo_filenames = np.array(photo_filenames)
    labels = np.load(dataset_dir+'/tag_train.npy')
    filenames_zip = list(zip(photo_filenames, labels))

    filenames_labels = []
    for image_str in filenames_zip:
        if image_str[0]!='898b451e5cb7169e074d2e893d7dd92bba70825a.jpg':
            filenames_labels.append(image_str)
    
    # 给分类名编号，并用字典形式表示
    #class_names_to_ids = dict(zip(class_names, range(len(class_names)))) # class_names和range(len(class_names))一一对应组成元组

    # Divide into train and test:
    random.seed(_RANDOM_SEED)
    random.shuffle(filenames_labels)

    data_root = os.path.join(dataset_dir, 'mydata')

    train_filenames_labels = filenames_labels[_NUM_VALIDATION:]
    train_zip_split = list(zip(*train_filenames_labels))
    train_squeeze_filenames = np.squeeze(train_zip_split[0])

    train_filenames = []
    for path in train_squeeze_filenames:
        _, filenames = os.path.split(path)
        train_filenames.append(os.path.join(data_root, filenames))

    train_labels = np.squeeze(train_zip_split[1])

    valid_filenames_labels = filenames_labels[:_NUM_VALIDATION]
    valid_zip_split = list(zip(*valid_filenames_labels))
    valid_squeeze_filenames = np.squeeze(valid_zip_split[0])

    valid_filenames = []
    for path in valid_squeeze_filenames:
        _, filenames = os.path.split(path)
        valid_filenames.append(os.path.join(data_root, filenames))

    valid_labels = np.squeeze(valid_zip_split[1])

    # First, convert the training and validation sets.
    _convert_dataset('train', train_filenames, train_labels,
                     dataset_dir)
    _convert_dataset('validation', valid_filenames, valid_labels,
                     dataset_dir)

    # Finally, write the labels file:
    # labels_to_class_names = dict(zip(range(len(class_names)), class_names))
    # dataset_utils.write_label_file(labels_to_class_names, dataset_dir)  # 根据labels_to_class_names在dataset_dir下编写labels.txt

    print('\nFinished converting the multilabels dataset!')
